#include "amba_net.h"

#include "utils.h"
AmbaNet::AmbaNet() noexcept : raw_{}, net_buf_{} {}

std::uint32_t AmbaNet::Init(std::string const &flexidag_path) noexcept {
  if (inited_) {
    return 0;
  }
  memset(&handle_, 0, sizeof(AMBA_CV_FLEXIDAG_HANDLE_s));
  memset(&init, 0, sizeof(AMBA_CV_FLEXIDAG_INIT_s));
  std::uint32_t Rval =
      Do_Load_And_Open_FlexidagBin(flexidag_path.c_str(), &net_buf_, &handle_);
  if (Rval != 0U) {
    Print("Error: loading and opening [%s] flexidag bin! Quit...\n",
          flexidag_path.c_str());
    return 1;
  }
  Rval = Do_Init_Flexidag(&handle_, &init);
  if (Rval != 0U) {
    Print("Error: init [%s] flexidag bin! Quit...\n", flexidag_path.c_str());
    return 1;
  }
  // Print("Do_Init_Flexidag success...\n");
  // Rval = alloc_IO_buffer();
  // if (Rval != 0U) {
  //   Print("Error encounted when handling input file! Quit...\n");
  //   return 1;
  // }
  // Print("alloc_IO_buffer success...\n");
  inited_ = true;
  return Rval;
}

std::uint32_t AmbaNet::RunBlock() noexcept {
  std::uint32_t flag = input_done();
  if (flag != 0U) {
    return flag;
  }
  Print("input_done success...\n");
  std::uint32_t Rval =
      AmbaCV_FlexidagRun(&handle_, &in_buf_, &out_buf_, &run_info_);
  for (int i = 0; i < out_buf_.num_of_buf; i++) {
    Rval = AmbaCV_UtilityCmaMemInvalid(&out_buf_.buf[i]);
  }
  return Rval;
}

// typedef void (*flexidag_cb)(void *vpHandle, uint32_t flexidag_output_num,
//                              flexidag_memblk_t *pblk_Output, void
//                              *vpParameter);

std::uint32_t AmbaNet::RunNonBlock(flexidag_cb call_back, void *cb_param,
                                   std::uint32_t &token_id) noexcept {
  std::uint32_t flag = input_done();
  if (flag != 0U) {
    return flag;
  }
  Print("input_done success...\n");
  return AmbaCV_FlexidagRunNonBlock(&handle_, call_back, cb_param, &in_buf_,
                                    &out_buf_, &token_id);
}

std::uint32_t AmbaNet::WaitRunFinish(std::uint32_t token_id) noexcept {
  return AmbaCV_FlexidagWaitRunFinish(&handle_, token_id, &run_info_);
}

// 从文件加载
std::uint32_t AmbaNet::LoadInPutFile(
    std::string const &input_path, int order,
    memio_source_recv_multi_raw_t *raw,
    flexidag_memblk_t *input_image_buf) noexcept {
  const char *input_image_path = input_path.c_str();
  Print("Loading input image from : %s\n", input_image_path);
  uint32_t size_align;
  uint32_t Rval = 0U;
  Rval = AmbaCV_UtilityFileSize(input_image_path, &size_align);
  if (Rval == 0U) {
    Rval = AmbaCV_UtilityCmaMemAlloc(size_align, 1, input_image_buf);
    if (Rval == 0U) {
      Rval = AmbaCV_UtilityFileLoad(input_image_path, input_image_buf);
    }
  }
  if (Rval != 0U) {
    Print("LoadInPutFile : load input image file fail (0x%x)\n", Rval);
  }

  if (Rval == 0U) {
    raw->num_io += 1;
    raw->io[order].addr = input_image_buf->buffer_caddr;
    raw->io[order].size = input_image_buf->buffer_size;
    raw->io[order].pitch = 0U;
  }
  return Rval;
}

/**
 * @brief 从 buffer 拷贝数据到到网络
 * @param [in] input_buffer         存放上一段网络输出数据的buffer
 * @param [in] buffer_size          数据大小，根据网络结构预先定好
 * @param [in] order                io顺序,要根据第二段网络的输入来确定
 * @param [in] raw                  io
 * @param [in] input_image_buf      memory block
 */
std::uint32_t AmbaNet::LoadInPutBuffer(
    char *input_buffer, uint32_t buffer_size, int order,
    memio_source_recv_multi_raw_t *raw,
    flexidag_memblk_t *input_image_buf) noexcept {
  Print("Loading input from buffer...\n");
  uint32_t Rval = 0U;
  Rval = AmbaCV_UtilityCmaMemAlloc(buffer_size, 1, input_image_buf);  // 4k 对齐
  if (Rval == 0U) {
    memcpy(input_image_buf->pBuffer, input_buffer, buffer_size);
    AmbaCV_UtilityCmaMemClean(input_image_buf);
  }

  if (Rval != 0U) {
    Print("LoadInPutFile : load input image file fail (0x%x)\n", Rval);
  }

  if (Rval == 0U) {
    raw->num_io += 1;
    raw->io[order].addr = input_image_buf->buffer_caddr;
    raw->io[order].size = input_image_buf->buffer_size;
    raw->io[order].pitch = 0U;
  }
  return Rval;
}

AMBA_CV_FLEXIDAG_IO_s &AmbaNet::GetOutBuf() noexcept { return out_buf_; }

// IO
std::uint32_t AmbaNet::alloc_IO_buffer() noexcept {
  uint32_t Rval = 0U;
  in_buf_.num_of_buf = raw_.num_io;
  Print("raw_.num_io = [%d] \n", raw_.num_io);
  for (uint32_t i = 0; i < in_buf_.num_of_buf; i++) {
    Rval = AmbaCV_UtilityCmaMemAlloc(sizeof(memio_source_recv_raw_t), 1,
                                     &in_buf_.buf[i]);  // 4k 对齐

    // memio_source_recv_raw_t * raw_ptr = (memio_source_recv_raw_t
    // *)in_buf_.buf[i].pBuffer; raw_ptr->addr = 1; raw_ptr->size = 2;

    // if (Rval == 0U) {
    //   memio_source_recv_raw_t *pDataIn;
    //   memcpy(&pDataIn, &in_buf_.buf[i].pBuffer, sizeof(void *));
    //   memcpy(pDataIn, &raw_.io[i], sizeof(memio_source_recv_raw_t));
    //   Rval = AmbaCV_UtilityCmaMemClean(&in_buf_.buf[i]);  // 写入数据
    // }
    if (Rval != 0U) {
      Print("Do_Alloc_IO_Buffer : in_buf.buf[%d] fail (0x%x)\n", i, Rval);
      break;
    }
  }

  if (Rval == 0U) {    
    out_buf_.num_of_buf = handle_.mem_req.flexidag_num_outputs;
    Print("out_buf_.num_of_buf = [%d] \n", out_buf_.num_of_buf);
    for (uint32_t i = 0; i < out_buf_.num_of_buf; i++) {
      Rval = AmbaCV_UtilityCmaMemAlloc(
          handle_.mem_req.flexidag_output_buffer_size[i], 1, &out_buf_.buf[i]);
      if (Rval != 0U) {
        Print("Do_Alloc_IO_Buffer : out_buf.buf[%d] fail (0x%x)\n", i, Rval);
        break;
      }
      // bev_net 这里把 15 个输出进行了 combine 1-13   14-15;
    }
  }

  return Rval;
}

std::uint32_t AmbaNet::input_done() noexcept {
  std::uint32_t flag{};
  for (uint32_t i = 0; i < in_buf_.num_of_buf; i++) {
    flag |= AmbaCV_UtilityCmaMemClean(&in_buf_.buf[i]);  // 写入数据
  }
  isAllocated = true;
  return flag;
}

std::uint32_t AmbaNet::destory() noexcept {
  std::uint32_t Rval = 0U;
  if (inited_) {
    Rval = AmbaCV_FlexidagClose(&handle_);
    Rval |= AmbaCV_UtilityCmaMemFree(&net_buf_);
    Rval |= AmbaCV_UtilityCmaMemFree(&init.state_buf);
    if (isAllocated) {
      Rval |= free_all_buffer();
      for (std::uint32_t i = 0; i < in_buf_.num_of_buf; i++) {
        Rval |= AmbaCV_UtilityCmaMemFree(&in_buf_.buf[i]);
      }
      for (std::uint32_t i = 0; i < out_buf_.num_of_buf; i++) {
        Rval |= AmbaCV_UtilityCmaMemFree(&out_buf_.buf[i]);
      }

      if (Rval != 0U) {
        Print("FREE_IO_Buffer fail (0x%x)\n", Rval);
      }
      isAllocated = false;
    }

    inited_ = false;
  }
  return Rval;
}
